from WindPy import w
import pandas as pd
from comm_tools.config import Config
from comm_tools import date_tool
from comm_tools.database_mysql import open_mysql, load_to_MySQL_on_Cloud

def format_to_pd(list_data):
    code_list = ["trade_code", "sec_name", "ipo_date", "ipo_collection", "ipo_leadundr"]

    df = pd.DataFrame(list_data)

    df = df.T
    df.columns = code_list

    return df

def save_data(df_data, table_name):
    with open_mysql() as engine:
        load_to_MySQL_on_Cloud(df_data, engine, table_name)

def upload(end_date=date_tool.get_last_day_of_previous_month()):
    # 初始化Wind API
    if not w.isconnected():
        w.start()  # 默认命令超时时间为120秒，如需设置超时时间可以加入waitTime参数，例如waitTime=60,即设置命令超时时间为60秒

    """ step 1 retrieve all listed co. ID """
    data_ID = w.wset("listedsecuritygeneralview", "sectorId=a001010100000000;field=wind_code")
    if data_ID.ErrorCode == 0:
        # 打印结果
        print("查询成功，结果如下：")
        for item in data_ID.Data:
            print(item)
    else:
        print(f"查询失败，错误代码：{data_ID.ErrorCode}")

    list_ID = data_ID.Data[0]
    str_ID = ', '.join(list_ID)

    """ step 2 retrieve ipo_collection data """
    data_ipo = w.wss(
        str_ID,
        "trade_code,  sec_name,  ipo_date,  ipo_collection,  ipo_leadundr", "unit=1;")

    df_ipo = format_to_pd(data_ipo.Data)

    """ step 3 filter by date """
    df_ipo["ipo_date"] = df_ipo["ipo_date"].dt.date
    start_date = date_tool.get_begin_date(end_date)
    df_ipo_last_month = df_ipo[(df_ipo["ipo_date"] >= start_date) & (df_ipo["ipo_date"] <= end_date)].copy()

    """ step 4 unit from yuan to yi """
    df_ipo_last_month["ipo_collection"] /= 100000000

    """ step 5 expand ipo_leadundr to multiple rows"""
    df_ipo_last_month["ipo_leadundr"] = df_ipo_last_month["ipo_leadundr"].str.split(",") # from str to list
    df_ipo_last_month = df_ipo_last_month.explode("ipo_leadundr", ignore_index=True) # expand to multiple rows
    print(df_ipo_last_month)

    # print(df_ipo_last_month)
    """ step 6 upload to database """
    c = Config()
    save_data(df_ipo_last_month, c.table_listed_co)


